# 环境准备：安装必要库
# pip install tensorflow numpy pillow

import tensorflow as tf
from tensorflow.keras import layers, models
import numpy as np

# 1. 数据准备
(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.mnist.load_data()

# 数据预处理
def preprocess_data(images, labels):
    # 扩展维度并归一化
    images = np.expand_dims(images, axis=-1).astype("float32") / 255.0
    # 转换标签为one-hot编码
    labels = tf.keras.utils.to_categorical(labels, 10)
    return images, labels

train_images, train_labels = preprocess_data(train_images, train_labels)
test_images, test_labels = preprocess_data(test_images, test_labels)

# 2. 构建CNN模型
model = models.Sequential([
    layers.Conv2D(32, (3,3), activation='relu', input_shape=(28,28,1)),
    layers.MaxPooling2D((2,2)),
    layers.Conv2D(64, (3,3), activation='relu'),
    layers.MaxPooling2D((2,2)),
    layers.Flatten(),
    layers.Dense(128, activation='relu'),
    layers.Dropout(0.5),
    layers.Dense(10, activation='softmax')
])

# 3. 编译模型
model.compile(optimizer='adam',
              loss='categorical_crossentropy',
              metrics=['accuracy'])

# 4. 训练模型
history = model.fit(train_images, train_labels,
                    epochs=10,
                    batch_size=128,
                    validation_split=0.2)

# 5. 评估模型
test_loss, test_acc = model.evaluate(test_images, test_labels)
print(f"Test accuracy: {test_acc:.4f}")

# 6. 保存模型
model.save('mnist_cnn.h5')